%% Section 0: Description of the file
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%% Explanation of what the file does
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 
% In this file we solve the RUV nested model using an algorithm
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%% Input files needed
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 
% InputData.xlsx
% ProcessedData.xlsx
%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%% Matlab functions invoked
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% 
% NEBA_Block1
% NEBA_Block2
% NEBA_Block3
% NECO_Block1
% NECO_Block2
% NECO_Block3

%% Section 1: Setup

clear all
clc

for caltype=1:4

    for flexindicator=0:1

        clearvars -except caltype flexindicator
        
        % Basic variables to be defined
        
        S             = 14;                 % Number of market sectors
        T             = 100;                % Number of periods (years)
        I             = 87;                 % Total number of regions
        M             = 50;                 % Number of US states
        sigma         = 6;                  % Trade elasticity
        chinapos      = 57;                 % Position of China in region ordering
        
        % Variables that might change depending on the calibration
        
        if caltype == 1 % Calibration 1: BASELINE
            mi        = 1;
            optdelta  = 4;
            delta     = 0.990721;
            ka        = 6.547537;
            nu        = 0.537175;
            tectime   = [0.646317,0.621760,1.632136,1.398806,0.942922,0.790867,...
                         0.862154]';
            tecsect   = [0.139465,0.077186,0.077833,0.053073,0.100555,0.070690,...
                         0.070298,0.081616,0.091727,0.083907,0.074990,0.078659]';
            ba        = 0.95;
            zrs       = 0.5;
            AECR      = [0,0,0,0,0];      
            AECRS     = [0,0,0,0,0];
            aeatw     = 1;
        elseif caltype == 2 % Calibration 2: Longer shock
            mi        = 1;
            optdelta  = 4;
            delta     = 0.993644;
            ka        = 13.531615;
            nu        = 0.705524;
            tectime   = [0.513774,0.573404,1.536869,1.373928,0.932361,0.779078,...
                         0.842860,0.582015,-0.46467,0.951536,0.879846]';
            tecsect   = [0.133088,0.075512,0.074893,0.088981,0.097700,0.071929,...
                         0.069008,0.074812,0.089183,0.080080,0.077719,0.067095]';
            ba        = 0.95;
            zrs       = 0.5;
            AECR      = [0,0,0,0,0];      
            AECRS     = [0,0,0,0,0];
            aeatw     = 1;
        elseif caltype == 3 % Calibration 3: No migration
            mi        = 0;
            optdelta  = 4;
            delta     = 0.990206;
            ka        = 6;
            nu        = 0.610981;
            tectime   = [0.646038,0.641380,1.651788,1.412809,0.952523,0.799161,...
                         0.871133]';
            tecsect   = [0.138887,0.077050,0.077822,0.053108,0.100538,0.070848,...
                         0.070252,0.081806,0.091811,0.084302,0.074913,0.078663]';
            ba        = 0.95;
            zrs       = 0.5;
            AECR      = [0,0,0,0,0];      
            AECRS     = [0,0,0,0,0];
            aeatw     = 1;
        elseif caltype == 4 % Calibration 4: nu=kappa
            mi        = 1;
            optdelta  = 4;
            delta     = 0.990919;
            ka        = 0.606216;
            nu        = 0.606216;
            tectime   = [0.645587,0.640116,1.648226,1.409473,0.949516,0.796212,...
                         0.867974]';
            tecsect   = [0.138960,0.077127,0.077794,0.053149,0.100670,0.070790,...
                         0.070231,0.081739,0.091817,0.084227,0.074958,0.078538]';
            ba        = 0.95;
            zrs       = 0.5;
            AECR      = [0,0,0,0,0];      
            AECRS     = [0,0,0,0,0];
            aeatw     = 1;
        end
        
        if flexindicator == 1
            delta     = 0.01;
        end
        
        % Import data
        
        dataimp       = xlsread('Inputs/ProcessedData.xlsx','BIFIXED','C2:CK1219');
        datarea       = reshape(dataimp',[I,I,S]);
        labshares     = xlsread('Inputs/InputData.xlsx','VA','C2:P88');
        inoutcdp      = xlsread('Inputs/InputData.xlsx','IO','D2:Q533');
        inouttemp1    = permute(reshape(inoutcdp',[S,S,I-M+1]),[2,1,3]);
        inouttemp2    = NaN(S,S,I);
        inouttemp2(:,:,1:M) = repmat(inouttemp1(:,:,1),1,1,M);
        inouttemp2(:,:,M+1:end) = inouttemp1(:,:,2:end);
        inouttemp3    = permute(inouttemp2,[3,1,2]);
        inoutmat      = inouttemp3 .* repmat(permute(1-labshares,[1,3,2]),1,S,1);
        
        % Define trade shares, labor income, deficits, final cons. shares, etc
        
        tradesharedat = datarea./repmat(sum(datarea,1),I,1,1);
        expbysec      = squeeze(sum(datarea,1));
        revbysec      = squeeze(sum(datarea,2));
        laborincsec   = labshares .* revbysec;
        deficits      = sum(expbysec,2)-sum(revbysec,2);
        alphasinter   = repmat(permute(revbysec,[1,3,2]),1,S,1);
        alphasregnum  = expbysec - sum(inoutmat .* alphasinter,3);
        Amat          = alphasregnum./repmat(sum(alphasregnum,2),1,S);
        Smat          = ones(1,S)*sigma;
        DeltaMat      = 0.01*ones(I,S);
        if optdelta == 1
            DeltaMat  = delta*ones(I,S);
        elseif optdelta == 2
            DeltaMat(:,1:12) = delta*ones(I,12);
        elseif optdelta == 3
            DeltaMat(1:M,:) = delta*ones(M,S);
        elseif optdelta == 4
            DeltaMat(1:M,1:12) = delta*ones(M,12);
        end
        
        % Shocks
        
        DefEra        = repmat(deficits,1,T);
        TECmatC       = ones(I,S,T);
        TAUmatC       = ones(I,I,S,T);
        
        % Generate matrix of mobility for all countries (irrelevant for USA)
        
        Splus         = S+1;
        EllMinus      = xlsread('Inputs/ProcessedData.xlsx','L1999MIX','B2:P88');
        EllInit       = xlsread('Inputs/InputData.xlsx','L2000CENSUS','B2:P88');
        EllResh       = reshape(EllInit',I*Splus,1)';
        EllDen        = repmat(sum(EllMinus,2),1,Splus)';
        EllDiv        = EllResh./EllDen(:)';
        EllDivE       = repmat(EllDiv,I*Splus,1);
        MuInter       = kron(eye(I),ones(Splus));
        MuInit        = MuInter.*EllDivE;
        
        % Assign USA matrices
        
        if mi == 1
            MuUSA     = xlsread('Inputs/ProcessedData.xlsx','MUFIXED','D2:ABY751');
        elseif mi == 0
            MuUSA     = xlsread('Inputs/ProcessedData.xlsx','MUFIXEDNM','D2:ABY751');
        elseif mi == 3
            UEllInit  = EllInit(1:M,:);
            UEllResh  = reshape(UEllInit',M*Splus,1)';
            UEllDiv   = UEllResh./sum(UEllInit(:));
            MuUSA     = repmat(UEllDiv,M*Splus,1);
        end
        if mi~=2
            MuInit(1:M*Splus,1:M*Splus) = MuUSA;
        elseif mi == 2
            MuInit    = eye(I*Splus);
        end
        Ugu           = ones(I,Splus,T);
        
        %% Section 2: Solve for the baseline economy
        
        lamdo         = min(min(nu,ka),0.6);
        iterations    = 0;
        distancebg    = 10;
        maxiter       = 1500;
        tolerance     = 3e-3;
        
        % Run the contraction mapping loop part of the tatonnement algorithm 
        % described in the paper
        
        while distancebg>tolerance && iterations<maxiter
            [MLBas,lLBas] = NEBA_Block1(Ugu,MuInit,EllInit,ba,ka,nu);
            
            [WDBas,LDBas,PDBas,RLBas,TLBas,YLBas] = NEBA_Block2(DefEra,...
            laborincsec,revbysec,tradesharedat,TAUmatC,TECmatC,labshares,...
            inoutmat,Smat,Amat,lLBas,DeltaMat,0.2,1e-4,AECR,AECRS,aeatw,M);
        
            UNew = NEBA_Block3(PDBas,Amat,WDBas,LDBas,lLBas,MLBas,ba,ka,nu,Ugu,zrs);
            
            iterations = iterations+1;
            distancebg = max(abs(UNew(:)-Ugu(:)));
            Ugu = lamdo * UNew + (1-lamdo)*Ugu;
            
            disp(['Distance = ',num2str(distancebg,'%.7f')])
        end
        
        %% Section 3: Solve for the counterfactual economy
        
        iterations    = 0;
        distancebg    = 10;
        Ugu2          = ones(I,Splus,T);
        lecs          = length(tectime);
        TECN          = ones(I,S,T);
        tecvalpre     = (1+tectime*tecsect').^(sigma-1);
        TECN(57,1:12,1:lecs) = permute(tecvalpre,[3,2,1]);
        TecHat        = TECN./TECmatC;
        
        % Run the contraction mapping loop part of the tatonnement algorithm 
        % described in the paper
        
        while distancebg>tolerance && iterations<maxiter
            [MLCfa,lLCfa,TZ] = NECO_Block1(Ugu2,MLBas,EllInit,ba,ka,nu);
            
            [WHCfa,LHCfa,PHCfa,RLCfa,TLCfa,YLCfa] = NECO_Block2(DefEra,...
            laborincsec,revbysec,tradesharedat,TAUmatC,TecHat,labshares,...
            inoutmat,Smat,Amat,lLCfa,DeltaMat,TLBas,WDBas,LDBas,0.3,1e-4,...
            AECR,AECRS,aeatw,M);
        
            [UNew2,Delrshat] = NECO_Block3(PHCfa,Amat,WHCfa,LHCfa,lLCfa,MLCfa,...
            ba,ka,nu,MLBas,TZ,lLBas,Ugu2,zrs,LDBas);
            
            iterations = iterations+1;
            distancebg = max(abs(UNew2(:)-Ugu2(:)));
            Ugu2 = lamdo * UNew2 + (1-lamdo) * Ugu2; 
            
            disp(['Distance = ',num2str(distancebg,'%.7f')])
        end
        
        %% Section 4: Compute welfare change (equivalent consumption variation)
        
        Pagg               = repmat(prod(PHCfa .^ repmat(Amat,1,1,T-1),2),1,S,1);
        EllRel             = lLCfa(:,2:end,:);
        EllBR              = lLBas(:,2:end,:);
        EllDot             = EllRel(:,:,2:end)./EllRel(:,:,1:end-1);
        EllDBas            = EllBR(:,:,2:end)./EllBR(:,:,1:end-1);
        EllHat             = EllDot./EllDBas;
        Omega              = ones(I,Splus,T-1);
        Omega(:,2:end,:)   = (WHCfa.*LHCfa)./(Pagg.*EllHat.*Delrshat);
        
        % Compute mu dots in baseline
        
        MuHashB            = NaN(I*Splus,I,T);
        MuCondB            = NaN(size(MLBas));
        MuHashC            = NaN(I*Splus,I,T);
        MuCondC            = NaN(size(MLCfa));
        for t=1:T
            MuHashB(:,:,t) = MLBas(:,:,t) * kron(eye(I),ones(Splus,1));
            MuCondB(:,:,t) = MLBas(:,:,t) ./ kron(MuHashB(:,:,t),ones(1,Splus));
            MuHashC(:,:,t) = MLCfa(:,:,t) * kron(eye(I),ones(Splus,1));
            MuCondC(:,:,t) = MLCfa(:,:,t) ./ kron(MuHashC(:,:,t),ones(1,Splus));
        end
        MuCondB(isnan(MuCondB)) = 0;
        MuCondC(isnan(MuCondC)) = 0;
        MuHashDB           = MuHashB(:,:,2:end)./MuHashB(:,:,1:end-1);
        MuHashDB(isnan(MuHashDB)) = 1;
        MuCondDB           = MuCondB(:,:,2:end)./MuCondB(:,:,1:end-1);
        MuCondDB(isnan(MuCondDB)) = 1;
        MuHashDC           = MuHashC(:,:,2:end)./MuHashC(:,:,1:end-1);
        MuHashDC(isnan(MuHashDC)) = 1;
        MuCondDC           = MuCondC(:,:,2:end)./MuCondC(:,:,1:end-1);
        MuCondDC(isnan(MuCondDC)) = 1;
        
        % Compute mu hats
        
        MuCondHat          = MuCondDC./MuCondDB;
        MuCondT1           = NaN(I*Splus,1,T-1);
        for counter=1:I*Splus
            MuCondT1(counter,1,:) = MuCondHat(counter,counter,:);
        end
        MuCondT2           = permute(reshape(MuCondT1,Splus,I,T-1),[2,1,3]);
        
        MuHashHat          = MuHashDC./MuHashDB;
        MuHashT1           = NaN(I*Splus,1,T-1);
        for counter=1:I*Splus
            MuHashT1(counter,1,:) = MuHashHat(counter,ceil(counter/Splus),:);
        end
        MuHashT2           = permute(reshape(MuHashT1,Splus,I,T-1),[2,1,3]);
        
        % Finally compute welfare change
        
        WELF1              = log(Omega./( MuCondT2.^nu .* MuHashT2.^ka));
        WELF2              = linspace(1,T-1,T-1)';
        WELF3              = permute(WELF2,[2,3,1]);
        WELF4              = repmat(WELF3,I,Splus,1);
        WELF5              = ba.^WELF4 .* WELF1;
        WELF               = sum(WELF5,3);
        
        % Change in chinese imports
        
        TOTInc             = repmat(sum(YLBas,2),1,S,1);
        consumption        = (TOTInc + repmat(deficits,1,S,T)) .* repmat(Amat,1,1,T);
        intermediat        = squeeze(sum(repmat(inoutmat,1,1,1,T) .* repmat(permute(RLBas,[1,4,2,3]),1,S,1,1),3));
        imfromchina        = squeeze(TLBas(57,:,:,:)) .* ( consumption + intermediat );
        imchinaus          = squeeze(sum(imfromchina(1:M,:,:),1));
        changeimpchina     = imchinaus(:,2:lecs+1)-imchinaus(:,1:lecs);
        changetimbas       = sum(changeimpchina(1:12,:),1)';
        changesecbas       = sum(changeimpchina,2);
        TOTInc             = repmat(sum(YLCfa,2),1,S,1);
        consumption        = (TOTInc + repmat(deficits,1,S,T)) .* repmat(Amat,1,1,T);
        intermediat        = squeeze(sum(repmat(inoutmat,1,1,1,T) .* repmat(permute(RLCfa,[1,4,2,3]),1,S,1,1),3));
        imfromchina        = squeeze(TLCfa(57,:,:,:)) .* ( consumption + intermediat );
        imchinaus          = squeeze(sum(imfromchina(1:M,:,:),1));
        changeimpchina     = imchinaus(:,2:lecs+1)-imchinaus(:,1:lecs);
        changetimcfa       = sum(changeimpchina(1:12,:),1)';
        changeseccfa       = sum(changeimpchina,2);
        changetimcha       = changetimcfa - changetimbas;
        changeseccha       = changeseccfa - changesecbas;
        
        name = strcat('Results',num2str(caltype));
        
        if flexindicator==1
            name = strcat('Results',num2str(caltype),'Flex');
        end
        
        parsave(name,changetimcha,changeseccha,AECRS,Amat,DefEra,Delrshat,...
        DeltaMat,EllBR,EllDBas,EllDen,EllDiv,EllDot,EllHat,EllInit,EllMinus,...
        EllRel,EllResh,I,LDBas,LHCfa,M,Omega,PDBas,PHCfa,Pagg,RLBas,RLCfa,S,...
        Smat,Splus,T,TECN,TECmatC,TecHat,UNew,UNew2,Ugu,Ugu2,WDBas,WELF,WELF1,...
        WELF2,WELF3,WELF4,WELF5,WHCfa,YLBas,YLCfa,aeatw,alphasinter,...
        alphasregnum,ba,chinapos,dataimp,datarea,deficits,delta,distancebg,...
        expbysec,inoutcdp,inoutmat,ka,lLBas,lLCfa,laborincsec,labshares,lecs,mi,...
        nu,optdelta,revbysec,sigma,tecsect,tectime,tecvalpre,tolerance,...
        tradesharedat,zrs,caltype,flexindicator)

        disp(['Finished with calibration number ',num2str(caltype),' with flexibility indicator ', num2str(flexindicator)])

    end

end
